import numpy as np

from task import TASK


class LOCOMOTION(TASK):
    
    
    def __init__(self) -> None:
        super().__init__()
        
        
    def get_fitness(self, position, orientation):
        
        if np.abs(orientation[2] - 0) < np.pi / 12:
            self.__fitness = (position[0] ** 2 + position[1] ** 2) ** 0.5
        else:
            self.__fitness = -(np.abs(orientation[2] - 0))
        
        return self.__fitness
    
        
    def is_finished(self, position, orientation, target):
        
        if ((position[0] - target[0]) ** 2 + (position[1] - target[0]) ** 2) ** 0.5 <= 1:
            return True
        
        else:
            return False


    def save_network(self, index):
        
        if self.network is not None:
            np.save(f"networks/locomotion/{index}-{self.__fitness:.3f}", self.network)
